#!/usr/bin/env python

# Constants:
PACKAGE_NAME = 'annotations'
NODE_NAME = 'check_accuracy'

# -------
# Imports:
# -------
# ROS
import roslib; roslib.load_manifest(PACKAGE_NAME)
import rospy, cv, rosbag
from annotations.msg import AnnotationArray, Annotation
from argparse import ArgumentParser

# ----------------------
class CompareAnnotations:
# ----------------------
	# -------------------------------------------------------
	def __init__(self, annotationBagPath, groundTruthBagPath):
	# -------------------------------------------------------
		# Initialize the node
		rospy.init_node(NODE_NAME)

		# Load the ground truth bag
		print 'Loading ground truth bag: %s' % groundTruthBagPath
		self.groundTruthBag = rosbag.Bag(groundTruthBagPath, 'r')

		# Load the annotations bag
		print 'Loading annotation bag: %s' % annotationBagPath
		self.annotationsBag = rosbag.Bag(annotationBagPath, 'r')

	# -------
	# Private|
	# ---------------------------------------------------------------------------------
	def __processAnnotationArrayMessages(self, annotationArrayMsg, groundTruthArrayMsg):
	# ---------------------------------------------------------------------------------
		truePos = 0
		falsePos = 0
		falseNeg = 0		

		# Loop through the annotation array
		for annotation in annotationArrayMsg.annotations:
			# Compare the annotation to ALL of the ground truth rectangles
			foundMatch = False
			for groundTruth in groundTruthArrayMsg.annotations:
				foundMatch |= self.__compareRectanglesForDetection(annotation, groundTruth)
			
			if foundMatch:
				truePos += 1
			else:
				falsePos += 1

		falseNeg = len(groundTruthArrayMsg.annotations) - truePos 
		return (truePos, falsePos, falseNeg)	

	# ----------------------------------------------------------------
	def __compareRectanglesForDetection(self, annotation, groundTruth):
	# ----------------------------------------------------------------
		""" Compares two rectangles to determine whether or not there was a detection """

		# Calculate the intersection of the rectangles
		intersection = self.__intersect(annotation, groundTruth)

		# If the rectangles don't intersect then it's not a valid detection
		if intersection == None:
			return False
		
		# Otherwise compare the area of the intersection -- if the are is >= 50% of the
		# area of the ground truth then count it as a detection
		x,y,w,h = intersection
		area = w * h
		groundTruthArea = groundTruth.width * groundTruth.height
		
		return (float(area) / float(groundTruthArea)) >= 0.5

	# ---------------------------
	def __intersect(self, r1, r2):
	# ---------------------------
		""" Calculate the intersection of two rectangles """
		tl = (max(r1.x, r2.x), max(r1.y, r2.y))
		br = (min(r1.x + r1.width, r2.x + r2.width), min(r1.y + r1.height, r2.y + r2.height))

		# Make sure the intersection is valid
		if tl[0] > br[0] or tl[1] > br[1]: 
			return None
		else:
			return (tl[0], tl[1], br[0] - tl[0], br[1] - tl[1])

	# ----------------------------------
	def __getTopicsByDataType(self, bag): 
	# ----------------------------------
	# Copied from rxbag (just to make sure we have access to it)
		""" 
		Get all the message types in the bag and their associated topics.
		
		@param bag: bag file
		@type  bag: rosbag.Bag
		@return: mapping from message typename to list of topics
		@rtype:  dict of str to list of str
		"""
		topics_by_datatype = {}
		for c in bag._get_connections():
			topics_by_datatype.setdefault(c.datatype, []).append(c.topic)
		
		return topics_by_datatype

	# ------
	# Public|
	# -----------
	def run(self):
	# -----------
		got = 0
		noGot = 0

		# Loop through and load the annotation messages
		# ---------------------------------------------
		self.annotations = {}
		annotationTopics =  self.__getTopicsByDataType(self.annotationsBag)
		annotationTopic = annotationTopics['annotations/AnnotationArray']
		annotationsGen = self.annotationsBag.read_messages(topics=annotationTopic)
		print 'Annotations:'
		for msgTup in annotationsGen:
			topic, msg, t = msgTup
			self.annotations[msg.header.stamp] = msg
		
		print len(self.annotations)

		# Loop through the ground truth messages
		# --------------------------------------
		groundTruthTopics = self.__getTopicsByDataType(self.groundTruthBag)
		groundTruthTopic = groundTruthTopics['annotations/AnnotationArray']
		groundTruthGen = self.groundTruthBag.read_messages(topics=groundTruthTopic)
		print 'Ground Truth:'
		for msgTup in groundTruthGen:
			topic, msg, t = msgTup
			if msg.header.stamp in self.annotations:
				got += 1
			else:
				noGot += 1

		print got
		print noGot

# -----------------------
if __name__ == '__main__':
# -----------------------
	# Get the path to the annotation & ground truth bags from the command line arguments	
	parser = ArgumentParser(description='Compare annotations to ground truth')
	parser.add_argument('-g', '--ground_truth', dest='ground_truth', help='Relative path to the ground truth bag')
	parser.add_argument('-a', '--annotations', dest='annotations', help='Relative path to the annotations bag')
	args = parser.parse_args()

	# Create the node
	node = None
	try:
		node = CompareAnnotations(args.annotations, args.ground_truth)
		node.run()
	except IOError as e:
		print 'Error opening bag file: %s' % str(e)
